#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''建造量化数据列表的工具
'''

import sys
import argparse
import random
from pathlib import Path

def build_args_parser():
    parser = argparse.ArgumentParser(prog = 'build_quant_list.py'
            , description = '根据目录内容生成文件列表')

    parser.add_argument('dirs'
            , type = str
            , nargs = '+'
            , help = 'dir list')

    parser.add_argument('-n'
            , '--number'
            , type = int
            , default = 200
            , help = '总共需要多少文件')

    parser.add_argument('-o'
            , '--output'
            , type = str
            , default = 'list.txt' 
            , help = '输出文件名')

    parser.add_argument('-ns'
            , '--no-shuffle'
            , action = 'store_true'
            , help = '不shuffle文件列表')

    parser.add_argument('-cs'
            , '--call-stack'
            , action = 'store_true'
            , help = '打印调用栈')

    parser.add_argument('-a'
            , '--absolute'
            , action = 'store_true'
            , help = '使用绝对路径')

    parser.add_argument('-dl'
            , '--dir-label'
            , action = 'store_true'
            , help = '文件所在的目录名作为label，以空格为分隔符，放在文件名后面')

    return parser

def build_file_item(args, f):
    if args.dir_label:
        # 文件所在目录就是label
        label = f.parent.parts[-1]
        return f'{f} {label}'
    else:
        return str(f)

def build_one_dir_list(args, files, p):
    if p.is_file():
        files.append(build_file_item(args, p))
    elif p.is_dir():
        for d in p.iterdir():
            # 递归调用
            build_one_dir_list(args, files, d)

def build_file_list(args):
    files = []
    for d in args.dirs:
        dp = Path(d)
        if not dp.exists():
            raise Exception(f'{dp} not found') 
        if not dp.is_dir():
            raise Exception(f'{dp} is not directory') 
        build_one_dir_list(args, files, dp)

    return files

def main(args):
    fileList = build_file_list(args)
    if not args.no_shuffle:
        random.shuffle(fileList)

    with open(args.output, 'w') as f:
        cnt = args.number if args.number < len(fileList) else len(fileList)
        for l in fileList[:cnt]:
            if args.absolute:
                f.write(f'{Path(l).absolute()}\n')
            else:
                f.write(f'{l}\n')

    print(f'{cnt} line wrote to {args.output}')

if __name__ == '__main__':
    try:
        parser = build_args_parser()
        args = parser.parse_args()

        main(args)
    except Exception as e:
        if args.call_stack:
            import traceback
            traceback.print_exc()
        else:
            print (f'*QuantException*: {e}')
        sys.exit(1)
